'''
Author: tomwoo tom.woo@outlook.com
Date: 2025-07-16 01:09:13
LastEditors: tomwoo tom.woo@outlook.com
LastEditTime: 2025-07-25 15:18:27
FilePath: /multi-modal_agents/main.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''

# 加载环境变量
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())

import os
from pathlib import Path
import shutil

import gradio as gr
from gradio_pdf import PDF

from ocr_pipelines import load_and_split_pdf_document, outputs_dir
from vss_agent import VssAgent, STATUS_OK


# Get environment variable with default value
vss_host = os.getenv('VSS_HOST')
vss_port = os.getenv('VSS_PORT')

pdf_path = "./pdf_docs"
pdf_file = "a-fabric-datasheet.pdf"
image_path = "./images"
image_1_file = "summarization_prompts_diagram.png"
image_2_file = "VSS_CA-RAG.png"
video_path = "./videos"
video_file = "warehouse.mp4"
prompt_example_path = "./prompt_examples"
prompt_example_file = "prompt_example.txt"
caption_summarization_prompt_example_file = "caption_summarization_prompt_example.txt"
summary_aggregation_prompt_example_file = "summary_aggregation_prompt_example.txt"
question_examples_path = "./question_examples"
question_examples_file = "question_examples.txt"


def cleanup_folder(folder_path):
    path = Path(folder_path)
    if path.exists():
        for item in path.iterdir():
            if item.is_dir():
                shutil.rmtree(item)  # Recursively delete subdirectories
            else:
                item.unlink()  # Delete files
        print(f"Folder '{folder_path}' has been cleaned up.")
    else:
        print(f"Folder '{folder_path}' does not exist.")


if vss_host is not None and vss_port is not None:
    vss_agent_enabled = True
else:
    vss_agent_enabled = False

Path(outputs_dir).mkdir(parents=True, exist_ok=True)

with gr.Blocks() as ocr_pipelines_tab:
    pdf = PDF(os.path.join(pdf_path, pdf_file), label="Upload a PDF Document", interactive=True)
    with gr.Row(variant="compact", equal_height=False):
        gr.Markdown("# Choose a type: ", container=False)
        radio = gr.Radio(choices=["Text", "Tables", "Images"], value="Text", container=False, interactive=True, scale=1)
        button = gr.Button("Extract and Display PDF Content", variant="primary", interactive=True, scale=4)

    @gr.render(inputs=[pdf, radio], triggers=[button.click])
    def extract_and_display_pdf_content(filename, type):
        cleanup_folder(outputs_dir)

        elements = load_and_split_pdf_document(filename, type)
        if elements is None:
            return

        num_elements = len(elements)
        match type:
            case "Text":
                for i, text_element in enumerate(elements):
                    gr.Textbox(text_element, max_lines=5, 
                               label=f"Text Element {i + 1} of {num_elements}", show_label=True, 
                               interactive=False, show_copy_button=True)
            case "Tables":
                for i, table_element in enumerate(elements):
                    gr.HTML(table_element, 
                            label=f"Table Element {i + 1} of {num_elements}", show_label=True, 
                            container=True)
            case "Images":
                gr.Gallery(elements, columns=5, 
                           label=f"{num_elements} Image Elements", show_label=True, 
                           elem_id="gallery", height="auto", object_fit="contain", interactive=False)
            case _:
                pass

prompt_example = Path(os.path.join(prompt_example_path, prompt_example_file)).read_text()
caption_summarization_prompt_example = Path(os.path.join(prompt_example_path, caption_summarization_prompt_example_file)).read_text()
summary_aggregation_prompt_example = Path(os.path.join(prompt_example_path, summary_aggregation_prompt_example_file)).read_text()
question_examples = Path(os.path.join(question_examples_path, question_examples_file)).read_text().split("\n")

vss_agent = VssAgent(vss_host, vss_port)

with gr.Blocks() as vss_agent_tab:
    def video_cleared():
        vss_agent.clear()

        return gr.update(interactive=False)

    def upload_button_clicked(filename):
        ret = vss_agent.upload(filename)
        if STATUS_OK != ret[0]:
            print(ret[1])
            return gr.update(interactive=False)

        return gr.update(interactive=True)

    def summarize_button_clicked(prompt, caption_summarization_prompt, summary_aggregation_prompt):
        ret = vss_agent.summarize(prompt, caption_summarization_prompt, summary_aggregation_prompt)
        if STATUS_OK != ret[0]:
            print(ret[1])
            return ret[1]

        return ret[1]

    def chat_func(question, history):
        ret = vss_agent.qna(question)
        if STATUS_OK != ret[0]:
            print(ret[1])
            return ret[1]

        return ret[1]

    with gr.Row(variant="compact", equal_height=True):
        gr.Image(os.path.join(image_path, image_1_file), label="Summarization Prompts Diagram", show_label=True, interactive=False)
        with gr.Column(variant="compact"):
            video = gr.Video(os.path.join(video_path, video_file), label="Upload a Video", show_label=True, interactive=True, autoplay=True, loop=True, scale=9)
            upload_button = gr.Button("Upload Video to Agent", variant="primary", interactive=vss_agent_enabled, scale=1)
        with gr.Column(variant="compact"):
            prompt = gr.Textbox(prompt_example, lines=4, 
                       label="Prompt (in English)", show_label=True, 
                       interactive=True, show_copy_button=True)
            caption_summarization_prompt = gr.Textbox(caption_summarization_prompt_example, lines=4, 
                       label="Caption Summarization Prompt (in English)", show_label=True, 
                       interactive=True, show_copy_button=True)
            summary_aggregation_prompt = gr.Textbox(summary_aggregation_prompt_example, lines=4, 
                       label="Summary Aggregation Prompt (in English)", show_label=True, 
                       interactive=True, show_copy_button=True)
    with gr.Row(variant="compact", equal_height=True):
        gr.Image(os.path.join(image_path, image_2_file), label="Video Q&A by Vector-RAG and Graph-RAG", show_label=True, interactive=False)
        with gr.Column(variant="compact"):
            summary = gr.Textbox("Video Summary", max_lines=30, 
                   label="Video Summary (in English)", show_label=True, 
                   interactive=False, show_copy_button=True, scale=29)
            summarize_button = gr.Button("Summarize Video", variant="primary", interactive=False, scale=1)
        chat_interface = gr.ChatInterface(chat_func, 
                                          type="messages", 
                                          textbox=gr.Textbox(placeholder="Input your question...", submit_btn=True, stop_btn=True), 
                                          title="Video Q&A", 
                                          description="This task currently only supports English.", 
                                          theme="default", 
                                          examples=question_examples, 
                                          cache_examples=False, 
                                          save_history=False)

    video.clear(video_cleared, outputs=[summarize_button])
    upload_button.click(upload_button_clicked, inputs=[video], outputs=[summarize_button])
    summarize_button.click(summarize_button_clicked, 
                           inputs=[prompt, caption_summarization_prompt, summary_aggregation_prompt], 
                           outputs=[summary])

demo = gr.TabbedInterface([ocr_pipelines_tab, vss_agent_tab], 
                          ["OCR Pipelines", "Video Summarization and Q&A Agent"], 
                          title="Multi-Modal AI Agents")


if __name__ == "__main__":
    demo.queue().launch()

# end of file
